#include "postprocess.h"
#include "utils.h"

// std::vector<std::string> class_names = 
//             { "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
//             "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
//             "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
//             "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
//             "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
//             "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
//             "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
//             "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
//             "hair drier", "toothbrush" };
     
// std::unordered_map<int, std::string> select_names = 
//             {{0, "person"}, {1, "bicycle"}, {2, "car"},
//             {3, "motorcycle"} , {5, "bus"}, {7, "truck"}, {9, "traffic light"},
//             {15, "cat"}, {16, "dog"},{ 39, "bottle"}, {62, "tv"}
//             };

cv::Rect get_rect(cv::Mat& img, float bbox[4]) {
  float l, r, t, b;
  float r_w = kInputW / (img.cols * 1.0);
  float r_h = kInputH / (img.rows * 1.0);
  if (r_h > r_w) 
  {
    l = bbox[0] - bbox[2] / 2.f;
    r = bbox[0] + bbox[2] / 2.f;
    t = bbox[1] - bbox[3] / 2.f - (kInputH - r_w * img.rows) / 2;
    b = bbox[1] + bbox[3] / 2.f - (kInputH - r_w * img.rows) / 2;
    l = l / r_w;
    r = r / r_w;
    t = t / r_w;
    b = b / r_w;
  } 
  else 
  {
    l = bbox[0] - bbox[2] / 2.f - (kInputW - r_h * img.cols) / 2;
    r = bbox[0] + bbox[2] / 2.f - (kInputW - r_h * img.cols) / 2;
    t = bbox[1] - bbox[3] / 2.f;
    b = bbox[1] + bbox[3] / 2.f;
    l = l / r_h;
    r = r / r_h;
    t = t / r_h;
    b = b / r_h;
  }
  return cv::Rect(round(l), round(t), round(r - l), round(b - t));
}

static float iou(float lbox[4], float rbox[4]) {
  float interBox[] = {
    (std::max)(lbox[0] - lbox[2] / 2.f , rbox[0] - rbox[2] / 2.f), //left
    (std::min)(lbox[0] + lbox[2] / 2.f , rbox[0] + rbox[2] / 2.f), //right
    (std::max)(lbox[1] - lbox[3] / 2.f , rbox[1] - rbox[3] / 2.f), //top
    (std::min)(lbox[1] + lbox[3] / 2.f , rbox[1] + rbox[3] / 2.f), //bottom
  };

  if (interBox[2] > interBox[3] || interBox[0] > interBox[1])
    return 0.0f;

  float interBoxS = (interBox[1] - interBox[0])*(interBox[3] - interBox[2]);
  return interBoxS / (lbox[2] * lbox[3] + rbox[2] * rbox[3] - interBoxS);
}

static bool cmp(const Detection& a, const Detection& b) 
{
  return a.conf > b.conf;
}

// 这里的输出直接就是1000*38个字节数，也就是1000个框，38是一个框的字节数
void nms(std::vector<Detection>& res, float* output, float conf_thresh, float nms_thresh) {
  int det_size = sizeof(Detection) / sizeof(float);   // det_size = 38  
  std::map<float, std::vector<Detection>> m;

  for (int i = 0; i < output[0] && i < kMaxNumOutputBbox; i++)  // kMaxNumOutputBbox最大检测数量1000
  { 
    // 框的第5个float为conf_thresh
    if (output[1 + det_size * i + 4] <= conf_thresh) continue; // 小于conf_thresh就滤掉

    // 筛选
    if(select_names.find(output[1 + det_size * i + 5]) == select_names.end())
    {
      // std::cout << output[1 + det_size * i + 5] << std::endl;
      continue;
    }
  
    Detection det;
    memcpy(&det, &output[1 + det_size * i], det_size * sizeof(float));  // 大于的结果就用det接受
    if (m.count(det.class_id) == 0) m.emplace(det.class_id, std::vector<Detection>());
    m[det.class_id].emplace_back(det);  // m就是一张图剩下的框,一个带序号的map
  }

  // 重新选择最大score的检测框，放入到输出矩阵中，
  // 然后遍历剩余的框，并盘算与最大框的IOU，
  // 当IOU大于Th时，那么该框受到抑制，从输入框中删除。小于阈值则保留。
  // 从输入框中删除。小于阈值则保留。
  for (auto it = m.begin(); it != m.end(); it++) 
  {
    auto& dets = it->second;
    std::sort(dets.begin(), dets.end(), cmp);  // 按自信度排序
    for (size_t m = 0; m < dets.size(); ++m) 
    {
      auto& item = dets[m];  // 一张图中第m个框
      res.emplace_back(item);  // res接受需要的框
      
      for (size_t n = m + 1; n < dets.size(); ++n) 
      {
        if (iou(item.bbox, dets[n].bbox) > nms_thresh)  // 非极大抑制
        {
          dets.erase(dets.begin() + n);
          --n;
        }
      }
    }
  }
}


void draw_bbox(cv::Mat& img_, std::vector<Detection>& res) 
{
    for(auto it : res)
    {
      cv::Rect box = get_rect(img_, it.bbox);
      cv::rectangle(img_, box, cv::Scalar(0, 0, 255), 2);
      
      std::stringstream ss;
      ss << std::fixed << std::setprecision(2) << it.conf;
      std::string s = class_names[it.class_id] + " " + ss.str();

      // cv::putText(img,  s, cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2);

      auto font_face = cv::FONT_HERSHEY_DUPLEX;
      auto font_scale = 0.8;
      int thickness = 1;
      int baseline=0;
      auto s_size = cv::getTextSize(s, font_face, font_scale, thickness, &baseline);
      
      if( box.tl().y - s_size.height < 0)
      {
        cv::rectangle(img_,
            cv::Point(box.tl().x, box.tl().y),
            cv::Point(box.tl().x + s_size.width, box.tl().y + s_size.height+7),
            cv::Scalar(0, 0, 255), -1);
        cv::putText(img_, s, cv::Point(box.tl().x, box.tl().y + s_size.height),
                font_face , font_scale, cv::Scalar(255, 255, 255), thickness);

      }
      else
      {
        cv::rectangle(img_,
              cv::Point(box.tl().x, box.tl().y - s_size.height - 7),
              cv::Point(box.tl().x + s_size.width, box.tl().y),
              cv::Scalar(0, 0, 255), -1);
        cv::putText(img_, s, cv::Point(box.tl().x, box.tl().y - 5),
                  font_face , font_scale, cv::Scalar(255, 255, 255), thickness);

      }
  }
}

static cv::Rect get_downscale_rect(float bbox[4], float scale) {
  float left = bbox[0] - bbox[2] / 2;
  float top = bbox[1] - bbox[3] / 2;
  float right = bbox[0] + bbox[2] / 2;
  float bottom = bbox[1] + bbox[3] / 2;
  left /= scale;
  top /= scale;
  right /= scale;
  bottom /= scale;
  return cv::Rect(round(left), round(top), round(right - left), round(bottom - top));
}

std::vector<cv::Mat> process_mask(const float* proto, int proto_size, std::vector<Detection>& dets) {
  std::vector<cv::Mat> masks;
  for (size_t i = 0; i < dets.size(); i++) {
    cv::Mat mask_mat = cv::Mat::zeros(kInputH / 4, kInputW / 4, CV_32FC1);
    auto r = get_downscale_rect(dets[i].bbox, 4);
    for (int x = r.x; x < r.x + r.width; x++) {
      for (int y = r.y; y < r.y + r.height; y++) {
        float e = 0.0f;
        for (int j = 0; j < 32; j++) {
          e += dets[i].mask[j] * proto[j * proto_size / 32 + y * mask_mat.cols + x];
        }
        e = 1.0f / (1.0f + expf(-e));
        mask_mat.at<float>(y, x) = e;
      }
    }
    cv::resize(mask_mat, mask_mat, cv::Size(kInputW, kInputH));
    masks.push_back(mask_mat);
  }
  return masks;
}

cv::Mat scale_mask(cv::Mat mask, cv::Mat img) {
  int x, y, w, h;
  float r_w = kInputW / (img.cols * 1.0);
  float r_h = kInputH / (img.rows * 1.0);
  if (r_h > r_w) {
    w = kInputW;
    h = r_w * img.rows;
    x = 0;
    y = (kInputH - h) / 2;
  } else {
    w = r_h * img.cols;
    h = kInputH;
    x = (kInputW - w) / 2;
    y = 0;
  }
  cv::Rect r(x, y, w, h);
  cv::Mat res;
  cv::resize(mask(r), res, img.size());
  return res;
}


